#type vertex
#version 330 core
layout (location = 0) in vec3 aPos;
layout (location = 1) in vec2 aTex;

out vec3 FragPos;

void main() {
    FragPos = aPos;
    gl_Position = vec4(aPos, 1.0);
}

////////////////////////////////////////////////////////////////
#type fragment
#version 330 core
////// From Vertex Shader //////
in vec3 FragPos;
out vec4 FragColor;
////////////////////////////////

uniform vec2 _raster_size;

#define PI 					3.1415926
#define TWO_PI 				6.2831852
#define FOUR_PI 			12.566370
#define INV_PI 				0.3183099
#define INV_TWO_PI 			0.1591549
#define INV_FOUR_PI 		0.0795775
#define FLOAT_MAX 3.402823466e+38
#define EPSILON 1.0e-3

bool Float_Equal(in float a, in float b) {
    return abs(a-b) < EPSILON;
}

/**********************************/
/*random*/
// int rdCount=0;
// uniform float _rdSeed[4];
// float texMul;
// float divtexMul;
// float RandXY(float x, float y){
//      return fract(cos(x * (12.9898) + y * (4.1414)) * 43758.5453);
// }
// float Rand(){
//     float a = RandXY(texMul, _rdSeed[0]);
//     float b = RandXY(_rdSeed[1], divtexMul);
//     float c = RandXY(rdCount++, _rdSeed[2]);
//     float d = RandXY(_rdSeed[3] , a);
//     float e = RandXY(b, c);
//     float f = RandXY(d, e);

//     return f;
// }
int seed = 1;
int rand(void) { seed = seed*0x343fd+0x269ec3; return (seed>>16)&32767; }
float frand() { return float(rand())/32767.0; }
vec2 frand2() { return vec2(frand(), frand()); }
vec3 frand3() { return vec3(frand(), frand(), frand()); }
void srand( ivec2 p, int frame )
{
    int n = frame;
    n = (n<<13)^n; n=n*(n*n*15731+789221)+1376312589; // by Hugo Elias
    n += p.y;
    n = (n<<13)^n; n=n*(n*n*15731+789221)+1376312589;
    n += p.x;
    n = (n<<13)^n; n=n*(n*n*15731+789221)+1376312589;
    seed = n;
}
/**********************************/
////// Ray //////
struct Ray {
    vec3 origin;
    vec3 direct;
    float tmin, tmax;
};
Ray Ray_Cons(in vec3 org, in vec3 dir) {
    Ray ret;
    ret.origin = org;
    ret.direct = dir;
    ret.tmin = EPSILON; 
    ret.tmax = FLOAT_MAX;
    return ret;
}
vec3 Ray_To(in Ray r, float t) {
    return r.origin + t * r.direct;
}
Ray Ray_Trans(in Ray r, mat4 trans) {
    Ray newRay;
    newRay.origin = (trans * vec4(r.origin, 1.0)).xyz;
    newRay.direct = mat3(trans) * r.direct;
    newRay.tmin = r.tmin;
    newRay.tmax = r.tmax;
    return newRay;
}
///// Sample /////
struct CameraSample {
    vec2 uv;
    vec2 lens;
};
//////////////////
////// Camera (Perspective) //////
struct Camera {
    vec3 position, left_bottom, horizontal, vertical;
};
uniform Camera _camera;
uniform vec3 _front;
Ray Camera_GenRay(in CameraSample ca_sample) {
    vec3 org = _camera.position;
    vec3 dir = _camera.left_bottom
                + ca_sample.uv.x * _camera.horizontal
                + ca_sample.uv.y * _camera.vertical
                - org;
    Ray genRay = Ray_Cons(org, normalize(dir));
    return genRay;
}
////// BRDF //////
#define MAT_DIFFUSE 0
#define MAT_GLOSSY 1
#define MAT_LIGHT 2
struct Material {
    int type;
    vec3 albedo;

    float roughness;
    float intensity;
};
#define MATERIAL_MAX_NUM 20
uniform Material _materials[MATERIAL_MAX_NUM];

#define HEMISPHERE_SAMPLE 1
vec3 UniformHemisphereSample(in vec2 rd) {
    float phi = TWO_PI * rd.y;
    float cosTheta = 1.0 - rd.x;
    float sinTheta = sqrt(1.0 - cosTheta * cosTheta);
    return vec3(cos(phi) * sinTheta, sin(phi) * sinTheta, cosTheta);
}
vec3 CosineHemisphereSample(in vec2 rd) {
    float phi = TWO_PI * rd.y;
    float cosTheta = sqrt(1.0 - rd.x);
    float sinTheta = sqrt(1.0 - cosTheta * cosTheta);
    return vec3(cos(phi) * sinTheta, sin(phi) * sinTheta, cosTheta);
}

vec2 ConcentricDiskSample(in vec2 rd) {
    vec2 uOffset = 2.0 * rd - vec2(1, 1);
 
    // Handle degeneracy at the origin
    if (Float_Equal(uOffset.x, 0) && Float_Equal(uOffset.y, 0)) return vec2(0, 0);
 
    // Apply concentric mapping to point
    float theta, r;
    if ((uOffset.x) > abs(uOffset.y)) {
        r = uOffset.x;
        theta = FOUR_PI * (uOffset.y / uOffset.x);
    } else {
        r = uOffset.y;
        theta = TWO_PI - FOUR_PI * (uOffset.x / uOffset.y);
    }
    return r * vec2(cos(theta), sin(theta));
}

vec3 toWorld(in vec3 a, in vec3 N) {
    vec3 B, C;
    if (abs(N.x) > abs(N.y)) {
        float invLen = 1.0 / sqrt(N.x * N.x + N.z * N.z);
        C = vec3(N.z * invLen, 0.0, -N.x *invLen);
    }
    else {
        float invLen = 1.0 / sqrt(N.y * N.y + N.z * N.z);
        C = vec3(0.0, N.z * invLen, -N.y *invLen);
    }
    B = cross(C, N);
    return a.x * B + a.y * C + a.z * N;
}

float pdf_diffuse(in vec3 n, in vec3 wo){
#if HEMISPHERE_SAMPLE 
    return max(dot(n, wo), 0.0) * INV_PI;
#else
    return INV_TWO_PI;
#endif
}

vec3 evalue_diffuse(in vec3 wi, in vec3 wo, in vec3 n) {
    return dot(wo, n) > 0.0 ? vec3(INV_PI) : vec3(0.0);
}

void sample_diffuse(in vec3 n, inout vec3 wo, in vec2 rd_sample, inout float pdf) {
#if HEMISPHERE_SAMPLE 
    wo = toWorld(UniformHemisphereSample(rd_sample), n);
#else
    wo = toWorld(CosineHemisphereSample(rd_sample), n);
#endif
    pdf = pdf_diffuse(n, wo);
}

float pdf_glossy(in vec3 n, in vec3 wo, in vec3 real_wo) {
    return 0.0;
}

vec3 evalue_glossy(in vec3 wi, in vec3 wo, in vec3 n) {
    return vec3(0.0);
}

void sample_glossy(in vec3 n, in vec3 wi, inout vec3 wo, in vec2 rd_sample, in float roughness) {
    wo = reflect(wi, n);
    wo += toWorld(vec3(roughness * ConcentricDiskSample(rd_sample), 0.0), n);
}

//////////////////
////// Intersection //////
struct Intersection {
    float tHit;
    int mat_id;
    vec3 pHit, Ng, Ns;
};
//////////////////////////
//////////////////////////////////
////// Sphere //////
#define PRIMITIVE_SPHERE 0
#define PRIMITIVE_TRIANGLE 1
struct Primitive {
    int type;

    bool is_active;
    mat4 trans;
    mat4 trans_inv;
    int mat_id;
    
    vec3 va, vb, vc;
};
#define PRIMITIVE_MAX_NUM 20
uniform int _primitive_num;
uniform Primitive _primitives[PRIMITIVE_MAX_NUM];
// from pbrt
bool primitive_Intersect(in Primitive pri, in Ray inRay, inout Intersection isct) {
    if (pri.type == PRIMITIVE_SPHERE) {
        Ray r = Ray_Trans(inRay, pri.trans_inv);
        float a = dot(r.direct, r.direct);
        float b = 2.0 * dot(r.direct, r.origin);
        float c = dot(r.origin, r.origin) - 1.0;
        float delta = b * b - 4.0 * a * c;
        if (delta < 0.0) return false;
        float rootD = sqrt(delta);
        float q = b < 0.0 ? -0.5 * (b - rootD) : -0.5 * (b + rootD);
        float t0 = q / a;
        float t1 = c / q;
        if (t0 > t1) { float t_swap=t0; t0=t1; t1=t_swap; }
        if (t0 > r.tmax || t1 < r.tmin)return false;
        float tHit = t0;
        if (t0 < r.tmin) {
            tHit = t1;
            if(tHit > r.tmax) return false;
        }
        vec3 pHit = Ray_To(r, tHit);
        if (Float_Equal(pHit.x, 0.0) && Float_Equal(pHit.y, 0.0))pHit.x = 1e-5;
        vec3 Ng = dot(pHit, r.direct) < 0.0 ? pHit : -pHit;
        isct.tHit = tHit;
        isct.pHit = (pri.trans * vec4(pHit, 1.0)).xyz;
        isct.Ng = normalize(mat3(transpose(inverse(pri.trans))) * Ng);
        isct.Ns = isct.Ng;
        isct.mat_id = pri.mat_id;
        return true;
    } else if (pri.type == PRIMITIVE_TRIANGLE) {
        Ray r = Ray_Trans(inRay, pri.trans_inv);
        vec3 ab_ = pri.vb - pri.va;
        vec3 ac_ = pri.vc - pri.va;
        vec3 face_N = cross(ab_, ac_);
        vec3 oa = pri.va - r.origin;
        float dot_d_n = dot(r.direct, face_N);
        if (Float_Equal(dot_d_n, 0.0)) return false;
        float tHit = dot(oa, face_N) / dot_d_n;
        if (tHit < r.tmin || tHit > r.tmax) return false;
        vec3 pHit = Ray_To(r, tHit);
        vec3 ap_ = pHit - pri.va;
        vec3 bp_ = pHit - pri.vb;
        vec3 cp_ = pHit - pri.vc;
        vec3 bc_ = pri.vc - pri.vb;
        vec3 ca_ = -ac_;
        vec3 abp_ = cross(ap_,ab_);
        vec3 bcp_ = cross(bp_,bc_);
        vec3 cap_ = cross(cp_,ca_);
        if (dot(abp_, bcp_) * dot(bcp_, cap_) < 0.0 || dot(cap_, abp_) * dot(bcp_, cap_) < 0.0) return false;
        isct.tHit = tHit;
        isct.pHit = (pri.trans * vec4(pHit, 1.0)).xyz;
        isct.Ng = dot(face_N, r.direct) < 0.0 ? face_N : -face_N;
        isct.Ng = normalize(mat3(transpose(inverse(pri.trans))) * isct.Ng);
        isct.Ns = isct.Ng;
        isct.mat_id = pri.mat_id;
        return true;
    }
    return false;
}

//////////////////////
uniform float _p_RR;
uniform bool _skybox;

bool scene_Intersect(inout Ray ray, inout Intersection isct) {
    bool isHited = false;
    for (int i = 0; i < _primitive_num && i < PRIMITIVE_MAX_NUM; i++) {
        if (_primitives[i].is_active && primitive_Intersect(_primitives[i], ray, isct)) {
            isHited = true;
            ray.tmax = isct.tHit;
        }
    }
    return isHited;
}

vec3 pt(float times) {
    vec3 Li = vec3(0.0);
    vec3 throughput = vec3(1.0);
    vec3 L_dir = vec3(1.0);
    CameraSample ca_sample;
    ca_sample.uv = (vec2(gl_FragCoord.x+frand(),
                         gl_FragCoord.y+frand()))
                          / _raster_size;
    Ray ray = Camera_GenRay(ca_sample);

    int bounce = 0;
    while(bounce < 6) {
        bool isHited = false;
        Intersection isct;
        isHited = scene_Intersect(ray, isct);        
        if (isHited) {
            { // light sampling
                // TODO
            }
            { // brdf sampling
                if(frand() > _p_RR) {
                    break;
                }
                float brdf_pdf;
                vec3 wo;
                if (_materials[isct.mat_id].type == MAT_LIGHT) {
                    throughput *= _materials[isct.mat_id].albedo * _materials[isct.mat_id].intensity;
                    Li += throughput;
                    break;
                } else if (_materials[isct.mat_id].type == MAT_DIFFUSE) {
                    sample_diffuse(isct.Ns, wo, frand2(), brdf_pdf);
                    throughput *= _materials[isct.mat_id].albedo * evalue_diffuse(ray.direct, wo, isct.Ns) * dot(wo, isct.Ns) / brdf_pdf / _p_RR;
                    ray.origin = isct.pHit;
                    ray.direct = wo;
                    // ray.origin = ray.origin + EPSILON * 10.0 * ray.direct;
                    ray.tmax = FLOAT_MAX;
                } else if (_materials[isct.mat_id].type == MAT_GLOSSY) {
                    sample_glossy(isct.Ns, ray.direct, wo, frand2(), _materials[isct.mat_id].roughness);
                    ray.origin = isct.pHit;
                    ray.direct = wo;
                    // ray.origin = ray.origin + EPSILON * 10.0 * ray.direct;
                    ray.tmax = FLOAT_MAX;
                }
            }
        } else {
            float t = -ray.direct.y * 0.5 + 0.5;
            throughput *= _skybox ? (1 - t) * vec3(0.3, 0.5, 0.7) + t * vec3(1.0) : vec3(0.0);
            Li += throughput;
            break;
        }
        bounce++;
    }
    return Li;
}
uniform int _spp;
uniform int _frame;
void main() {
    srand(ivec2(gl_FragCoord), _frame);
    // texMul = gl_FragCoord.x + gl_FragCoord.y / gl_FragCoord.x;
	// divtexMul = - gl_FragCoord.x / gl_FragCoord.y + gl_FragCoord.y;
    vec3 color = vec3(0.0);
    for(int i = 0; i < _spp; i++) 
        color += pt(i);
    color /= _spp;
    FragColor = vec4(pow(color, vec3(1.0/2.2)),1.0);
} 
